#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
from huggingface_hub import snapshot_download

def download_model():
    """下载sentence-transformers模型到本地"""
    
    print("=== 模型下载器 ===")
    
    # 模型名称
    model_name = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
    
    print(f"📥 正在下载模型: {model_name}")
    print("⏳ 这可能需要几分钟时间，请耐心等待...")
    
    try:
        # 下载模型到本地缓存
        local_path = snapshot_download(
            repo_id=model_name,
            local_dir="./models/sentence-transformers",
            local_dir_use_symlinks=False
        )
        
        print(f"✅ 模型下载完成！")
        print(f"📁 本地路径: {local_path}")
        
        # 显示文件大小
        total_size = 0
        for root, dirs, files in os.walk(local_path):
            for file in files:
                file_path = os.path.join(root, file)
                total_size += os.path.getsize(file_path)
        
        print(f"📊 总大小: {total_size / 1024 / 1024:.2f} MB")
        
        return True
        
    except Exception as e:
        print(f"❌ 下载失败: {e}")
        return False

def test_model_loading():
    """测试模型加载"""
    print("\n=== 测试模型加载 ===")
    
    try:
        from sentence_transformers import SentenceTransformer
        
        # 尝试加载本地模型
        model_path = "./models/sentence-transformers"
        if os.path.exists(model_path):
            print(f"🔍 尝试从本地加载模型: {model_path}")
            model = SentenceTransformer(model_path)
            print("✅ 本地模型加载成功！")
            
            # 测试编码
            test_text = ["这是一个测试句子"]
            embeddings = model.encode(test_text)
            print(f"✅ 编码测试成功，维度: {embeddings.shape}")
            
            return True
        else:
            print("❌ 本地模型路径不存在")
            return False
            
    except Exception as e:
        print(f"❌ 模型加载失败: {e}")
        return False

def main():
    """主函数"""
    print("=" * 60)
    print("Sentence-Transformers 模型下载器")
    print("=" * 60)
    
    # 下载模型
    success = download_model()
    
    if success:
        # 测试加载
        test_model_loading()
    
    print("\n🎉 完成！")

if __name__ == "__main__":
    main() 